# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/10 18:31 
@Author  : ZhangShenao 
@File    : 3.使用configurable_alternatives方法动态替换运行时组件.py 
@Desc    :

使用configurable_alternatives方法,可以动态替换运行时组件
比如在构建好的Chain中,可以动态替换掉prompt、llm和parser等组件,而不是替换组件内的参数信息
"""
import os

import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import ConfigurableField
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()
prompt = ChatPromptTemplate.from_template('{query}')

# 创建llm组件,并定义可替换的备选组件
llm = ChatZhipuAI(model='glm-4-air').configurable_alternatives(
    # 定义可替换的组件
    ConfigurableField(id='llm', name='LLM大语言模型', description='可替换多种不同的语言模型'),
    # 定义默认组件key
    default_key='zhipu',
    # 设置可替换组件的key和value
    gpt4=ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE')),
)

parser = StrOutputParser()
chain = prompt | llm | parser

# 根据可替换的不同组件,创建两个chain
zhipu_chain = chain.with_config(configurable={
    'llm': 'zhipu'
})
gpt4_chain = chain.with_config(configurable={
    'llm': 'gpt4'
})

# 调用两个chain
query = '你好，你是哪个大模型？'
zhipu_result = zhipu_chain.invoke({'query': query})
gpt_result = gpt4_chain.invoke({'query': query})

print(f'zhipu_result: {zhipu_result}')
print(f'gpt_result: {gpt_result}')

# zhipu_result: 你好，我是基于智谱 AI 公司训练的大型语言模型开发的人工智能助手。我的任务是针对用户的问题和要求提供适当的答复和支持。我不是某个特定的大模型，而是根据具体的应用场景和需求，使用了不同的大型语言模型。
# gpt_result: 您好，我是GPT-3，一个由OpenAI开发的大型语言模型。有什么可以帮助您的吗？
